from output import Output
from plugin_system import Plugin
from parameter_system import Enum

class XYZOutput(Output):
    def __init__(self, filename=None, freq=None):
        """
        Takes the filename of the file that the data
        will be written to, the frequency at which the
        data will be written
        """
        self.filename = filename
        self.fl = None
        self.frequency = freq

    def initialize(self):
        """
        Opens the output file.
        """
        self.fl = open(self.filename, "w")

    def finalize(self):
        """
        Closes the output file.
        """
        if self.fl == None:
            raise Exception, "Cannot finalize output that has not be initialized."
        self.fl.close()

    def output(self, system, step):
        if step % self.frequency == 0:
            # we expect a list of 3-tuples
            num_objects = len(system.planet_names)
            header = "%s %s\n" % (num_objects, step)
            self.fl.write(header)
            for name in system.planet_names:
                pos = system.positions[name]
                vel = system.velocities[name]
                s = "%s %s %s %s %s %s %s\n" % (name, pos[0], pos[1], pos[2], \
                                                       vel[0], vel[1], vel[2])
                self.fl.write(s)
            

class XYZOutputPluginAdapter(Plugin, XYZOutput):
    plugin_name = "XYZOutput"
    parameter_definitions = {"filename" : (str, True, "File to which to write data."),
                             "freq" : (int, True, "Frequency (in steps) at which to write data.")}
    def __init__(self, *args, **user_params):
        params = self.configure_parameters(**user_params)
        XYZOutput.__init__(self, filename=params["filename"], freq=params["freq"])
